#!/usr/local/bin/python
#-*- coding: UTF-8 -*-
'''
Created on 2021年3月25日

@author: root
'''
import seaborn as sns
#import util.dbcom as dbcom
#import util.mdbcom as mdbcom
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import plotly.plotly
import plotly.graph_objs as graphObj
import dataMan as DM
import util.xtool as xtool
import util.tradeDate as TD
import apps.adog.util.dbcom as dbcom
import apps.adog.util.mdbcom as mdbcom
#from apps.adog.util.dbcom import replaceManyToMysql

#SELECT concat('"',cod,'":"',min(date),'"') FROM future.jq_hq_1d where code_type='fin8' group by cod
'''
'2021-02-02', 'PK'
'2021-01-11', 'LH'
'2020-11-20', 'BC'
'2020-10-13', 'PF'
'2020-06-23', 'LU'
'2020-03-31', 'PG'
'2019-12-09', 'SA'
'2019-09-26', 'EB'
'2019-09-25', 'SS'
'2019-08-16', 'RR'
'2019-08-12', 'NR'
'2019-08-09', 'UR'
'2019-04-30', 'CJ'
'2018-12-10', 'EG'
'2018-11-27', 'SP'
'2018-08-17', 'TS'
'2018-03-26', 'SC'
'2017-12-22', 'AP'
'2017-08-18', 'CY'
'2016-01-04', 'BU'
'2016-01-04', 'C'
'2016-01-04', 'CF'
'2016-01-04', 'CS'
'2016-01-04', 'CU'
'2016-01-04', 'FB'
'2016-01-04', 'FG'
'2016-01-04', 'FU'
'2016-01-04', 'HC'
'2016-01-04', 'I'
'2016-01-04', 'IC'
'2016-01-04', 'IF'
'2016-01-04', 'IH'
'2016-01-04', 'J'
'2016-01-04', 'JD'
'2016-01-04', 'JM'
'2016-01-04', 'JR'
'2016-01-04', 'L'
'2016-01-04', 'LR'
'2016-01-04', 'M'
'2016-01-04', 'MA'
'2016-01-04', 'NI'
'2016-01-04', 'OI'
'2016-01-04', 'P'
'2016-01-04', 'PB'
'2016-01-04', 'PM'
'2016-01-04', 'PP'
'2016-01-04', 'RB'
'2016-01-04', 'RI'
'2016-01-04', 'RM'
'2016-01-04', 'RS'
'2016-01-04', 'RU'
'2016-01-04', 'SF'
'2016-01-04', 'SM'
'2016-01-04', 'SN'
'2016-01-04', 'SR'
'2016-01-04', 'T'
'2016-01-04', 'TA'
'2016-01-04', 'TC'
'2016-01-04', 'TF'
'2016-01-04', 'V'
'2016-01-04', 'WH'
'2016-01-04', 'WR'
'2016-01-04', 'Y'
'2016-01-04', 'ZC'
'2016-01-04', 'ZN'
'2013-12-06', 'BB'
'2012-05-10', 'AG'
'2008-01-09', 'AU'
'2005-01-04', 'A'
'2005-01-04', 'AL'
'2005-01-04', 'B'
'''

class CorrMan():
    def __init__(self):
        self.data=None
        self.fi=None
    
        
    def loadData(self,dataMan=None,cod=[],startDate=None,endDate=None,freq='1d'):
        dm=None
        if dataMan is None:
            dm=DM.DataMan(startDate=startDate,endDate=endDate,obs=cod,freq=freq)
        else:
            dm=dataMan
        self.data=dm.data
        self.fi=dm.future_info
        
    def getCorr(self,startDate=None,endDate=None,obs=[],fields=['date','cod','chg_cls_r']):
        df=self.data.copy()
        df=df.loc[(df['date']<=endDate)&(df['date']>=startDate),fields]
        if len(obs)>0:
            df=df.loc[df['cod'].isin(obs),:]
        df.sort_values(by='cod',ascending=1,inplace=True)
        df.drop_duplicates(['date','cod'], keep='last', inplace=True)
        df.reset_index(drop=True,inplace=True)
        pivot_data=df.pivot(index='date',columns='cod')[fields[-1]]
        corr=pivot_data.astype(float).fillna(0.0).corr() 
        return pivot_data,corr
        
    def setCorr(self,startDate=None,endDate=None,obs=[],fields=['date','cod','chg_cls_r']):
        self.df=self.df.loc[(self.df['date']<=endDate)&(self.df['date']>=startDate),fields]
        if len(obs)>0:
            self.df=self.df.loc[self.df['cod'].isin(obs),:]
        self.df.sort_values(by='cod',ascending=1,inplace=True)
        self.df.drop_duplicates(['date','cod'], keep='last', inplace=True)
        self.df.reset_index(drop=True,inplace=True)
        data=self.df.pivot(index='date',columns='cod')[fields[-1]]
        self.corr=data.astype(float).fillna(0.0).corr()
        
        #return self.corr
        #corr.sort_values(by='A',ascending=1,inplace=True)
        #self.corr.to_csv('corr.csv')
    
    
  
    
    
        
    def drawPlot(self):
        #plotly
        x=self.corr.columns.values.astype(str)
        y=self.corr.columns.values.astype(str)
        z=self.corr.values.astype(float)
                
        trace = graphObj.Heatmap(
                           z=z,
                           x=x,
                           y=y,
                           #autocolorscale=True
                           hoverinfo='x+y+z',
                           )
        layout = graphObj.Layout(
                    title='',
                    annotations=graphObj.Annotations([
                                graphObj.Annotation(
                                    x=0.1,
                                    y=1.,
                                    showarrow=False,
                                    text='',
                                    xref='paper',
                                    yref='paper'
                                )
                            ]),
                    autosize=False,
                    width=1600,
                    height=1600,
                    margin=dict(
                        l=5,
                        r=5,
                        b=40,
                        t=40
                    )
                )
                               
        fig = graphObj.Figure(data=[trace],layout=layout)
        #plotly.offline.plot(fig,auto_open=False,filename='heatMap.html') 
        output=plotly.offline.plot(fig, output_type='div',include_plotlyjs=False)
        return output
    
    def drawSns(self):
        sns.pairplot(self.corr) 
        sns.clustermap(self.corr)
        plt.show()
        #sns_plot = sns.pairplot(df, hue='species', size=2.5)
        #fig = sns_plot.get_figure()
        #return fig
    
if __name__=='__main__':
    
    #cm=CorrMan(obs=['A','L','PP','V','M'])
    cm=CorrMan()
    #cm.drawPlot()
    #cm.drawSns()
    startDate='2017-01-01'
    endDate=xtool.nowDate()
    tradeDate=TD.tradeDateDf
    tradeDate=tradeDate.loc[(tradeDate['tradeDate']>=startDate)&(tradeDate['tradeDate']<endDate)]
    tradeDate.reset_index(drop=True, inplace=True)
    cm.loadData(startDate=startDate, endDate=endDate,freq='1d')
    
   
    total_index=TD.getTradeDayDiff(startDate,endDate )
    
    for i in range(50):
        print(i)
        period=np.random.randint(80, high=300,size=1)[0]
        start_index=np.random.randint(0, high=tradeDate.index.max()-period,size=1)
        end_index=start_index+period
        
        sDate=tradeDate.loc[start_index,'tradeDate'].values[0]
        eDate=tradeDate.loc[end_index,'tradeDate'].values[0]
        
        _,corr=cm.getCorr(startDate=sDate,endDate=eDate)
        corr=corr.stack()
        corr.index.names=['cod1','cod2']
        corr=corr.reset_index(drop=False)
        corr.rename(columns={0:'corr'},inplace=True)
        corr['startDate']=sDate
        corr['endDate']=eDate
        corr['period']=period
        corr['round']=i
        corr['_id']=str(i)+'_'+corr['cod1']+'_'+corr['cod2']
       
        mdbcom.saveBatch('corr_test', corr.to_dict(orient='records'))
        dbcom.creatTableReplaceMany('corr_test', corr, drop_exist=False)
        #dbcom.replaceManyToMysql('corr_test', corr.to_dict(orient='records'))
    
    